from mcp.server.fastmcp import FastMCP
from typing import Union, List, Dict, Any
import math

# 创建MCP服务实例
mcp = FastMCP("Advanced Math Operations Service")

@mcp.tool()
def calculate(operation: str, a: Union[int, float], b: Union[int, float]) -> str:
    """基本数学运算工具
    
    Args:
        operation: 运算类型，支持 'add' (加), 'subtract' (减), 'multiply' (乘), 'divide' (除)
        a: 第一个操作数
        b: 第二个操作数
        
    Returns:
        运算结果或错误信息
    """
    operations = {
        'add': lambda x, y: x + y,
        'subtract': lambda x, y: x - y,
        'multiply': lambda x, y: x * y,
        'divide': lambda x, y: x / y if y != 0 else "错误：除零错误"
    }
    
    if operation not in operations:
        return f"不支持的运算类型：{operation}。支持的运算：{', '.join(operations.keys())}"
    
    try:
        result = operations[operation](a, b)
        return f"{a} {operation} {b} = {result}"
    except Exception as e:
        return f"运算失败：{str(e)}"

@mcp.tool()
def complex_calculation(operation: str, real1: float, imag1: float, real2: float = 0, imag2: float = 0) -> str:
    """复数运算工具
    
    Args:
        operation: 运算类型，支持 'add' (加), 'subtract' (减), 'multiply' (乘), 'divide' (除)
        real1: 第一个复数的实部
        imag1: 第一个复数的虚部
        real2: 第二个复数的实部
        imag2: 第二个复数的虚部
        
    Returns:
        复数运算结果
    """
    # 构建复数
    complex1 = complex(real1, imag1)
    complex2 = complex(real2, imag2)
    
    operations = {
        'add': lambda x, y: x + y,
        'subtract': lambda x, y: x - y,
        'multiply': lambda x, y: x * y,
        'divide': lambda x, y: x / y if y != 0 else "错误：除零错误"
    }
    
    if operation not in operations:
        return f"不支持的运算类型：{operation}。支持的运算：{', '.join(operations.keys())}"
    
    try:
        result = operations[operation](complex1, complex2)
        return f"({real1} + {imag1}i) {operation} ({real2} + {imag2}i) = {result.real} + {result.imag}i"
    except Exception as e:
        return f"运算失败：{str(e)}"

@mcp.tool()
def statistics_calculation(data: List[float], operation: str) -> str:
    """统计计算工具
    
    Args:
        data: 数据列表
        operation: 运算类型，支持 'mean' (均值), 'median' (中位数), 'mode' (众数), 'std' (标准差), 'var' (方差)
        
    Returns:
        统计计算结果
    """
    if not data:
        return "错误：数据列表不能为空"
    
    operations = {
        'mean': lambda x: sum(x) / len(x),
        'median': lambda x: sorted(x)[len(x)//2] if len(x) % 2 != 0 else (sorted(x)[len(x)//2 - 1] + sorted(x)[len(x)//2]) / 2,
        'mode': lambda x: max(set(x), key=x.count),
        'std': lambda x: (sum((i - sum(x)/len(x))**2 for i in x) / len(x))**0.5,
        'var': lambda x: sum((i - sum(x)/len(x))**2 for i in x) / len(x)
    }
    
    if operation not in operations:
        return f"不支持的运算类型：{operation}。支持的运算：{', '.join(operations.keys())}"
    
    try:
        result = operations[operation](data)
        return f"{operation} of {data} = {result}"
    except Exception as e:
        return f"运算失败：{str(e)}"

@mcp.tool()
def solve_equation(equation_type: str, coefficients: Dict[str, float]) -> str:
    """方程求解工具
    
    Args:
        equation_type: 方程类型，支持 'linear' (线性方程), 'quadratic' (二次方程)
        coefficients: 方程系数
        
    Returns:
        方程的解
    """
    operations = {
        'linear': lambda coeffs: -coeffs.get('b', 0) / coeffs.get('a', 1) if coeffs.get('a', 1) != 0 else "错误：系数a不能为0" if 'a' in coeffs else "错误：缺少系数a",
        'quadratic': lambda coeffs: 
            "错误：缺少系数a" if 'a' not in coeffs else
            "错误：系数a不能为0" if coeffs['a'] == 0 else
            (lambda a, b, c: 
                f"x = {-b/(2*a)} (重根)" if b**2 - 4*a*c == 0 else
                f"x1 = {(-b + math.sqrt(b**2 - 4*a*c))/(2*a)}, x2 = {(-b - math.sqrt(b**2 - 4*a*c))/(2*a)}" if b**2 - 4*a*c > 0 else
                f"x1 = {(-b)/(2*a)} + {math.sqrt(4*a*c - b**2)/(2*a)}i, x2 = {(-b)/(2*a)} - {math.sqrt(4*a*c - b**2)/(2*a)}i")
            (coeffs['a'], coeffs.get('b', 0), coeffs.get('c', 0))
    }
    
    if equation_type not in operations:
        return f"不支持的方程类型：{equation_type}。支持的类型：{', '.join(operations.keys())}"
    
    try:
        result = operations[equation_type](coefficients)
        return f"{equation_type}方程的解: {result}"
    except Exception as e:
        return f"求解失败：{str(e)}"

# 启动服务
if __name__ == "__main__":
    mcp.run(transport="stdio")